import re
import ast
import json
from datetime import datetime
from src.common import commonUtils
from src.common.logger import getLogger
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import SystemMessage, HumanMessage

logger = getLogger()

class ReactAgent:

    def __init__(self, llm_model, agent_tools, max_steps: int = 3):
        self.llm_model = llm_model
        self.agent_tools = agent_tools
        self.max_steps = max_steps

    def parse_content_output(self, content):
        thought_match = re.search(r"思考：(.+?)\n", content, re.DOTALL)
        action_match = re.search(r"行动：(.+?)\n", content)
        action_input_match = re.search(r"行动输入：(.+?)\n", content)
        observation_match = re.search(r"观察：(.+?)\n", content)
        return {
            "thought": thought_match.group(1).strip() if thought_match else "",
            "action": ast.literal_eval(action_match.group(1).strip()) if action_match else "",
            "action_input": json.loads(action_input_match.group(1)) if action_input_match else "",
            "observation": observation_match.group(1).strip() if observation_match else "",
            "answer": content.split("答案：")[-1].strip()
        }

    def invoke(self, query):
        logger.info(f"ReactAgent invoke query: {query}")
        tool_names = list(self.agent_tools.keys())
        tool_descs = commonUtils.build_tools_description(self.agent_tools)

        history = []
        for step in range(self.max_steps):
            logger.info(f"ReactAgent invoke step: {step}")
            template = """
                你是一个智能AI助手，使用ReAct框架逐步解决问题，必须使用中文回答。
                
                当前时间是：{current_date}
                可用的工具：{tool_descs}
                当前任务：{input}
                
                请严格按以下格式和要求回答：
                思考：[你的推理过程]，请列出详细推理步骤
                行动：{tool_names}中的工具名，如果有多个工具，以列表格式返回，如果没有行动，返回为空
                行动输入：[工具的输入参数]，以json格式返回，如果没有行动输入，返回为空
                观察：[工具返回的结果]
                ...（重复思考/行动/观察）
                思考：严格根据上下文和观察进行思考推理是否能得出最终答案
                答案：[问题的最终答案]，严格根据上下文严谨的回答，如果没有答案或不知道，返回为空
            """
            current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            prompt = template.format(input = query, current_date = current_date, tool_names = tool_names, tool_descs = tool_descs)

            response = (self.llm_model | StrOutputParser()).invoke([
                SystemMessage(content = "你是一个严谨的问题助手"),
                HumanMessage(content = prompt)
            ])
            logger.info(f"ReactAgent invoke response: \n{response}")

            parsed_content = self.parse_content_output(response)
            logger.info(f"ReactAgent invoke parsed_content: {parsed_content}")
            history.append(parsed_content)
            if parsed_content["answer"]:
                return parsed_content["answer"]

            actions = parsed_content["action"]
            logger.info(f"ReactAgent invoke actions: {actions}")
            if actions:
                for action in actions:
                    if action not in self.agent_tools:
                        raise Exception(f"无效工具名：{action}")

                    tool = self.agent_tools[action]
                    if tool:
                        observation = tool.run(parsed_content["action_input"])
                        logger.info(f"ReactAgent invoke action: {action}, action_input: {parsed_content["action_input"]}, result len: {len(observation)} ")
                        history[-1]["observation"] = observation
                        query += f"\n观察：{observation}"
        return f"达到最大步骤数 {self.max_steps} 仍未找到答案"
